from jose import jwt
from fastapi.security import OAuth2PasswordBearer
from fastapi import HTTPException, status, Request
from server.report.models.model import User
from .mysql_utils import mysql_util
from typing import Union
from datetime import datetime, timedelta
    

class MyOAuth2PasswordBearer(OAuth2PasswordBearer):
    def __init__(self, tokenUrl: str):
        super().__init__(
            tokenUrl=tokenUrl,
            scheme_name="",
            scopes=None,
            description="JWT Authentication",
            auto_error=True
        )
        self.SECRET_KEY = "ed970259a19edfedf1010199c7002d183bd15bcaec612481b29bac1cb83d8137"
        self.ALGORITHM = "HS256"
        self.ACCESS_TOKEN_EXPIRE_MINUTES = 30

        self.credentials_exception = HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Could not validate credentials",
            headers={"WWW-Authenticate": "Bearer"},
        )
         
    def create_jwt_token(self, data: dict, expires_delta: Union[timedelta, None] = None):
        # 设置 token 过期时间
        if expires_delta:
            expire = datetime.utcnow() + expires_delta
            data.update({"exp": expire})
        
        # 进行jwt加密
        token = jwt.encode(claims=data, key=self.SECRET_KEY, algorithm=self.ALGORITHM)
        return token

    def check_jwt_token(self, token):
        try:
            pyload = jwt.decode(
                token,
                self.SECRET_KEY,
                algorithms=[self.ALGORITHM]
            )

            user_name: str = pyload.get("user_name")

            with mysql_util.get_db() as db:
                self.user = db.query(User).filter_by(user_name=user_name).first()

            if not self.user:
                raise self.credentials_exception
            return self.user
        except (jwt.JWTError, jwt.ExpiredSignatureError, AttributeError):
            # 抛出自定义异常， 然后捕获统一响应
            raise self.credentials_exception 
         
    async def __call__(self, request: Request):

        try:

            path: str = request.get('path')

        except Exception as e:
            # 捕获异常并抛出自定义错误e
            raise HTTPException(status_code=401, detail="Invalid token or request")
        
        if path.endswith('/login') | path.endswith('/register') |  path.endswith('/docs') |path.endswith('/') | path.__contains__('/ws/msg'):
            return ""
        
        try:
            authorization = request.headers["Authorization"]

            scheme, _, token = authorization.partition(" ")

            if scheme.lower() != "bearer":
                token = authorization

        except Exception:
            raise self.credentials_exception
        
        self.user = self.check_jwt_token(token)

        return self.user
    

my_oauth2_scheme = MyOAuth2PasswordBearer(tokenUrl='user/login')